package main

import (
	"container/list"
	"fmt"
	"net"
	"sync"
	"time"
)

type SmartRateLimiter struct {
	maxRate     int        // 最大允许速率（包/秒）
	currentRate int        // 当前允许速率
	packetQueue *list.List // 缓存队列
	queueSize   int        // 最大队列长度
	lastAdjust  time.Time  // 最后调整时间
	mu          sync.Mutex
}

func NewSmartRateLimiter(maxRate, queueSize int) *SmartRateLimiter {
	return &SmartRateLimiter{
		maxRate:     maxRate,
		currentRate: maxRate / 2, // 初始设为最大值的一半
		packetQueue: list.New(),
		queueSize:   queueSize,
		lastAdjust:  time.Now(),
	}
}

func (s *SmartRateLimiter) Allow(packet []byte) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 如果队列未满，直接接收并放入队列
	if s.packetQueue.Len() < s.queueSize {
		s.packetQueue.PushBack(packet)
		return true
	}

	// 队列已满时，根据当前速率决定是否丢弃
	if s.packetQueue.Len() >= s.queueSize &&
		s.packetQueue.Len() > s.currentRate/2 {
		return false
	}

	s.packetQueue.PushBack(packet)
	return true
}

func (s *SmartRateLimiter) AdjustRate() {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	s.lastAdjust = now

	// 动态调整速率
	queueLen := s.packetQueue.Len()
	switch {
	case queueLen > s.queueSize*3/4:
		s.currentRate = int(float64(s.currentRate) * 0.9) // 队列快满，降低速率
	case queueLen < s.queueSize/4:
		s.currentRate = int(float64(s.currentRate) * 1.1) // 队列较空，提高速率
	}

	// 确保在合理范围内
	if s.currentRate < 1 {
		s.currentRate = 1
	}
	if s.currentRate > s.maxRate {
		s.currentRate = s.maxRate
	}
}

func (s *SmartRateLimiter) ProcessQueue(processFunc func([]byte)) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 根据当前速率处理队列中的包
	toProcess := s.currentRate / 10 // 每100ms处理一次
	if toProcess < 1 {
		toProcess = 1
	}

	for i := 0; i < toProcess && s.packetQueue.Len() > 0; i++ {
		e := s.packetQueue.Front()
		processFunc(e.Value.([]byte))
		s.packetQueue.Remove(e)
	}
}

func main() {
	limiter := NewSmartRateLimiter(1000, 5000) // 最大1000包/秒，队列5000包

	// 启动速率调整协程
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		for range ticker.C {
			limiter.AdjustRate()
		}
	}()

	// 启动处理协程
	go func() {
		ticker := time.NewTicker(100 * time.Millisecond)
		for range ticker.C {
			limiter.ProcessQueue(func(packet []byte) {
				// 实际处理业务逻辑
				fmt.Printf("Processing packet: %s\n", string(packet))
			})
		}
	}()

	// UDP服务器
	conn, err := net.ListenPacket("udp", ":8080")
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	buf := make([]byte, 1500)
	for {
		n, addr, err := conn.ReadFrom(buf)
		if err != nil {
			continue
		}

		packet := make([]byte, n)
		copy(packet, buf[:n])

		if !limiter.Allow(packet) {
			fmt.Printf("Rate limited (current rate: %d/s) - dropped packet from %s\n",
				limiter.currentRate, addr)
			continue
		}

		fmt.Printf("Accepted packet from %s (queue: %d)\n",
			addr, limiter.packetQueue.Len())
	}
}
